// Programmer: Aaron Steele
// Purpose: Calculate the codon frequencies (Refactoring of CodonMultipliers.pl)
// Usage : codonmultipliers <inputfile> <outputfile>

#include <stdio.h>
#include <unistd.h>
#include <ctype.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <errno.h>
#include <stdbool.h>

int unique_codons = 64;
int unique_aa = 64;

/*An array of all possible codons*/
char codon_types[64][4] = {"TTT", "TTC", "TTA","TTG", "CTT", "CTC", "CTA", "CTG",
						 "ATT", "ATC", "ATA", "ATG", "GTT", "GTC", "GTA", "GTG",
						 "TCT", "TCC", "TCA", "TCG", "AGT", "AGC", "CCT", "CCC",
						 "CCA", "CCG", "ACT", "ACC", "ACA", "ACG", "GCT", "GCC", 
						 "GCA", "GCG", "TAT", "TAC", "CAT", "CAC", "CAA", "CAG",
						 "AAT", "AAC", "AAA", "AAG", "GAT", "GAC", "GAA", "GAG",
						 "TGT", "TGC", "TGG", "CGT", "CGC", "CGA", "CGG", "AGA", 
						 "AGG", "GGT", "GGC", "GGA", "GGG", "TAA", "TAG", "TGA"};

/* Maps each codon in codon_types to the corresponding amino acid */
char aa_codon_map[64] = {'F', 'F', 'L', 'L', 'L', 'L', 'L', 'L', 'I', 'I', 'I',
						 'M', 'V', 'V', 'V', 'V', 'S', 'S', 'S', 'S', 'S', 'S',
						 'P', 'P', 'P', 'P', 'T', 'T', 'T', 'T', 'A', 'A', 'A',
						 'A', 'Y', 'Y', 'H', 'H', 'Q', 'Q', 'N', 'N', 'K', 'K',
						 'D', 'D', 'E', 'E', 'C', 'C', 'W', 'R', 'R', 'R', 'R',
						 'R', 'R', 'G', 'G', 'G', 'G', 'X', 'X', 'X'};
 
/*An array of all amino acids (x corresponds to stop codon)*/
char aa_types[21] = {'F','L','I', 'M', 'V', 'S', 'P', 'T', 'A', 'Y', 'H', 'Q',
					 'N', 'K', 'D', 'E', 'C', 'W', 'R', 'G', 'X'};

/* Translate a codon to the corresponding Amino Acid */
char codon2aa(char *codon){
	int i;
	for(i=0; i<unique_codons; ++i){
		if(strcmp(codon,codon_types[i]) == 0)
			return aa_codon_map[i];
	}
}

/* Get the index of a given codon from the codon_types array */
int get_codon_index(char *codon){
	int i;
	for(i=0; i<unique_codons; ++i){
		if(strcmp(codon,codon_types[i]) == 0){ 	
			return i;		
		}
	}
	return -1;
}

/* Get the index of the amino acid with respect to the aa_types */
int get_aa_index(char aa){
	int i;
	for(i=0; i<unique_aa; ++i){
		if(aa == aa_types[i])
			return i;
	}
	return -1;
}

/*Get the next sequence from a given input fasta file*/
int getSequence(FILE *in, char seq[]){
	int d,s, length = 0;
	char* descr[300];

	if(in == NULL){
		return -1;
	}
	
	//Read the description of a sequence which ends with \n
	d = fscanf(in, "%[^\n]%*c", descr);
	
	if( d == EOF)
		return -1;
	
	//Read the actual sequence which is terminated by the '>'
	s = fscanf(in, "%[^>]%n%*c", seq, &length); 

	if( s == EOF)
		return -1;

	//Remove \n from the end of the sequence
	if(seq[length-1] == '\n'){
		seq[length-1] = '\0';
		length = length -1;
	}

	//Make all the characters uppercase for quality assurance
	int i;
	for(i=0; i<length; ++i)
		seq[i] = toupper(seq[i]);

	return length;
}

//Check to see if given nucleotides are a valid codon
bool is_valid_codon(char* nucleotides){
	int index = get_codon_index(nucleotides);	
	if(index == -1){
		return false;
	}

	return true;
}

int main(int argc, char **argv){
	
	if(argc != 3){
		fprintf(stderr, "No input/output fasta file specified, or too many arguments\n");
		exit(1);
	}

	//Input-Output File variables
	FILE *infile;
	FILE *outfile;
	int index = 1;
	char *filename = NULL;

	//Multiplier variables
	int window = 17;
	int half = 8;
	int seq_count = 0;	//Total Sequences
	int bad_seq = 0;	//Sequences with incomplete reading frames
	int bad_codon = 0;	//Unrecognized codons (not ATGC)


	//Open the input file
	filename = argv[index];
	infile = fopen(filename, "rt");
	if(infile == NULL){
		fprintf(stderr, "Failed to open %s\n", filename);
		exit(1);
	}
	index++;

	//Open the output file
	filename = argv[index];
	outfile = fopen(filename, "wt");
	if(outfile == NULL){
		fprintf(stderr, "Failed to open %s\n", filename);
		exit(1);
	}
	
	/*Create a 2D array for counting codons and amino acids in their respective
 	 * locations in the window	*/
	int aa_counts[unique_aa][window];
	int codon_counts[unique_codons][window];


	//Read the sequences from the file
	char seq[40000];
	int seq_length= -1;
	int i,j,k;
	
	//Indicies of the codons in the window
	int center_codon;
	int last_codon;
	int cur_codon;
	bool valid;
	
	//Char buffers for holding codons
	char cdn1[3];
	char cdn2[3];

	//Index of the codons being compared (with repsect to amino acids)
	int center_aa_index;
	int center_cdn_index;
	int cur_aa_index;
	int cur_cdn_index;
		
	//Initialize Codon Distance Tracking Array
	// usage: [center_cdn_index][cur_cdn_index][distance]
	int cdn_dist[64][64][window];
	
	//Initialize AA Distance Tracking Array
	// usage: [center_cdn_index][cur_aa_index][distance]
	int aa_dist[64][21][window];
	
	//fprintf(stdout, "Setting all values in the distance array\n");
	//Set all values in the distance arrays to 0
	for(i=0; i<64; ++i){
		for(j=0; j<64; ++j){
			for(k=0; k<window; ++k){
				cdn_dist[i][j][k] = 0;
				if(j < 21)
					aa_dist[i][j][k]=0;
			}
		}
	}


	seq_length = getSequence(infile, seq);

//	fprintf(stdout, "Length: %d\n", seq_length);
//	fprintf(stdout, "Sequence: \n%s\n", seq);

	while(seq_length != -1){
		seq_count++;

		//Check the quality of the sequence
		if(seq_length % 3 != 0)
			bad_seq++;
		
		//If sequence is good use it
		else{
			int codon_count = seq_length/3;
			int distance = 0; 	//Distance between codons
			int array_distance = 0;
			char aa;
	
			//Extract and use each codon
			for(i=0; i<codon_count; ++i){
				center_codon = i;
				memcpy(cdn1,&seq[center_codon*3],3);
		
				//Check if codon is valid			
				valid = is_valid_codon(cdn1);
				if(valid == false){
					bad_codon++;
				}else{
					//Get the codons in the window
					cur_codon = center_codon - half;
					if(cur_codon < 0) { cur_codon = 0;}
				
					last_codon = center_codon + half;
					if(last_codon >= codon_count){last_codon = codon_count -1;}
				
					//Get the index of each codon
					aa = codon2aa(cdn1);
					center_cdn_index = get_codon_index(cdn1);
					center_aa_index = get_aa_index(aa);
							
					while(cur_codon < last_codon){
						memcpy(cdn2,&seq[cur_codon*3],3);
						cur_cdn_index = get_codon_index(cdn2);
						aa = codon2aa(cdn2);
						cur_aa_index = get_aa_index(aa);
 
						distance =	cur_codon - center_codon;						
						array_distance = distance + half;
												
						cdn_dist[center_cdn_index][cur_cdn_index][array_distance]++;
						aa_dist[center_cdn_index][cur_aa_index][array_distance]++;	
						
						cur_codon++;
					}

				}			
			}

		}
		//Read the next sequence from the file
		seq_length = getSequence(infile,seq);
	
	}


	//Calculate and Output the results to the file
		
	//Header
	fprintf(outfile, "\t\t");
	for(i=0; i<64; ++i)
		fprintf(outfile,"%c\t",aa_codon_map[i]);
	
	fprintf(outfile, "\n\t\t");
	for(i=0; i<64; ++i)
		fprintf(outfile, "%s\t", codon_types[i]);
	
	//fprintf(stdout, "Declaring calc vars\n");
	//Calculations
	int cdncount; // total ocrences of codon 2 within windows with codon 1 (except -1 and 0)
	int AAcountm1; //Total occrences of AA2 at -1 distance(before codon1)
	int cdncountm1; // Total occrences of codon1 at -1 distance (before codon2)
	int AAcount; // total occrences of AA2 within windows with codon 1 (except -1 and 0)
	int aa_index;

	//Result variables
	double ratio;
	double avefreq;
	double freqm1;

	for(i=0; i<64; ++i){
		fprintf(outfile, "\n%c\t",aa_codon_map[i]);
		fprintf(outfile, "%s\t", codon_types[i]);
		//fprintf(stdout, "Printing codon: %s\n", codon_types[i]);
	
		for(j=0; j<64; ++j){
			AAcount = 0;
			cdncount = 0;
			cdncountm1 = 0;
			AAcountm1 = 0;

			aa_index = get_aa_index(aa_codon_map[j]);
			//fprintf(stdout, "AAINDEX: %d\n", aa_index);

			for(k=0; k<window; ++k){
				if(k != half && k != (half-1)){
					cdncount += cdn_dist[i][j][k];
					AAcount += aa_dist[i][aa_index][k];
				}
				if(k == half-1){
					AAcountm1 += aa_dist[i][aa_index][k];
					cdncountm1 += cdn_dist[i][j][k];
				}			
			}
			
			//Print the result
			if(AAcount == 0 || AAcountm1 == 0)
				fprintf(outfile, "0.00\t");
			else{
				avefreq = cdncount/(double)AAcount;
				freqm1 = cdncountm1/(double)AAcountm1;
				
				if(avefreq > 0){
					ratio = freqm1/(double)avefreq;	
					fprintf(outfile, "%f\t",ratio);
				}
				else{fprintf(outfile,"0.00\t");}
			}			
		
		}	
	}	
	fprintf(outfile,"\n");

	fclose(infile);
	fclose(outfile);

	fprintf(stdout, "Total Sequences: %d\n", seq_count);
	fprintf(stdout, "Bad Sequences: %d\n", bad_seq);
	fprintf(stdout, "Unrecognized Codons: %d\n", bad_codon);
}
